from loguru import logger
from torchtext.data.utils import get_tokenizer

from config.config import cfg
from utils.data_helper import NewsDataset
from utils.utils import build_vocab
from utils.utils import load_dataset, log_init

if __name__ == '__main__':
    log_init('test_data_loader', log_dir=cfg.log_dir)

    # 创建英文token分词器，此处直接使用torchtext工具类提供的 get_tokenizer 方法获取
    tokenizer = get_tokenizer('basic_english')
    # 构建词汇表 【英文单词 <----> 数字索引】 映射表
    vocab = build_vocab(tokenizer=tokenizer, split='train', data_file_path=cfg.dataset_dir)

    logger.debug('-' * 30 + ' 测试：使用vocab词表构建文本内容 idx 索引向量测试 ' + '-' * 30)
    # 获取train.csv第一个样本
    text = "Reuters - Short-sellers, Wall Street's dwindling\band of ultra-cynics, are seeing green again."
    logger.debug(f'text: {text}')
    # 构建文本idx索引向量
    tokens = tokenizer(text)
    logger.debug(f'tokens: {tokens}')
    text_idx = vocab.lookup_indices(tokens)
    logger.debug(f'text_index: {text_idx}')

    logger.debug('-' * 30 + ' 测试：NewsDataset 工具类实现数据集train/valid拆分及batch批次封装 ' + '-' * 30)
    # 加载训练数据集
    train_data = load_dataset(cfg.dataset_dir, split='train')
    # 实例化NewsDataset对象
    new_dataset = NewsDataset(tokenizer=tokenizer, vocab=vocab, config=cfg)
    # 将数据集拆分为训练集与验证集合
    train_dataloader, valid_dataloader = new_dataset.train_valid_split(train_data, ratio=0.95)

    # 查看bath封装后数据
    for idx, (labels, texts, offsets) in enumerate(valid_dataloader):
        logger.debug(f'first batch labels: {labels}')
        logger.debug(f'first batch offsets: {offsets}')
        logger.debug(f'first batch text: {texts}')
        break
